import cv2
import math
import yanmo
import numpy as np
from time import sleep


class Calculate(yanmo.ColorDetector):  # 继承修改后的ColorDetector类

    def __init__(self, image, x1, y1):
        super().__init__(image)  # 调用父类初始化
        self.x1 = x1  # 目标点x坐标
        self.y1 = y1  # 目标点y坐标
        self.x, self.y, self.w, self.h = 0, 0, 0, 0  # 检测区域坐标

    def _calculate_angle(self, x, y, w, h):
        """计算检测区域中心与目标点的角度

        Args:
            x, y: 检测区域左上角坐标
            w, h: 检测区域宽高

        Returns:
            计算得到的角度值(500-2500)
        """
        x_center = x + w // 2
        y_center = y + h // 2
        dy = x_center - self.x1
        dx = y_center - self.y1
        angle = math.atan2(dy, dx)
        angle = math.degrees(angle)
        return int(2000 / 270 * angle + 500)

    def shibie(self, n):
        """识别图像中的颜色区域并计算角度

        Returns:
            tuple: (角度结果列表, 颜色mask, 标记后的图像)
        """
        mask = self.show_colors_combined(n)  # 使用父类的detect_all_colors方法
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        display_frame = self.image.copy()
        results = []

        for i, cnt in enumerate(contours, 1):
            self.x, self.y, self.w, self.h = cv2.boundingRect(cnt)
            if self.w < 50 or self.h < 50:  # 过滤小区域
                continue

            # 绘制检测区域和编号
            cv2.rectangle(display_frame, (self.x, self.y), (self.x + self.w, self.y + self.h), (0, 255, 0), 2)
            cv2.putText(display_frame, str(i), (self.x, self.y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

            # 计算角度并保存结果
            p = self._calculate_angle(self.x, self.y, self.w, self.h)
            results.append({p: (self.x, self.y)})
            print(f"角度{p}对应坐标: ({self.x}, {self.y})")

        return results, mask, display_frame


class ArmController(Calculate):
    def __init__(self, image, x1, y1):
        super().__init__(image, x1, y1)
        self.L1, self.L2, self.L3 = 134.5, 134.5, 103.5
        self.pixel_to_mm = 380.0 / (635 ** 2 + 490 ** 2) ** 0.5
        self.base_x, self.base_y = self.x1, self.y1

        print(f"\n机械臂参数:")
        print(f"物理长度: L1={self.L1}mm, L2={self.L2}mm, L3={self.L3}mm")
        print(f"像素比例: {self.pixel_to_mm:.4f} mm/pixel")
        print(f"基座位置: ({self.base_x}, {self.base_y})px")

    def image_to_arm_coords(self, x_pixel, y_pixel):
        dx = x_pixel - self.base_x
        dy = y_pixel - self.base_y
        return dx * self.pixel_to_mm, dy * self.pixel_to_mm

    def inverse_kinematics_all(self, x, y, num_samples=100):
        d = np.sqrt(x ** 2 + y ** 2)
        solutions = []

        # 可达性检查
        if not (abs(self.L1 - self.L2 - self.L3) <= d <= (self.L1 + self.L2 + self.L3)):
            print(f"目标距离{d:.1f}mm 超出工作空间[{abs(self.L1 - self.L2 - self.L3):.1f}-{self.L1 + self.L2 + self.L3:.1f}]mm")
            return solutions

        for theta3 in np.linspace(0, 2 * np.pi, num_samples):
            A = self.L2 + self.L3 * np.cos(theta3)
            B = self.L3 * np.sin(theta3)
            C = (x ** 2 + y ** 2 + self.L1 ** 2 - (A ** 2 + B ** 2)) / (2 * self.L1)
            r = np.sqrt(x ** 2 + y ** 2)

            if np.abs(C) > r:
                continue

            alpha = np.arctan2(y, x)
            gamma = np.arccos(np.clip(C / r, -1.0, 1.0))

            for theta1 in [alpha - gamma, alpha + gamma]:
                x_prime = x - self.L1 * np.cos(theta1)
                y_prime = y - self.L1 * np.sin(theta1)

                denom = A ** 2 + B ** 2
                if denom < 1e-6:
                    continue

                cos_theta12 = (A * x_prime + B * y_prime) / denom
                sin_theta12 = (A * y_prime - B * x_prime) / denom

                if np.abs(cos_theta12) > 1 or np.abs(sin_theta12) > 1:
                    continue

                theta12 = np.arctan2(sin_theta12, cos_theta12)
                theta2 = (theta12 - theta1)  # 俯视图不需要+π调整
                solutions.append([theta1, theta2, theta3])

        return solutions

    # 保持其他方法不变
    def angle_to_pulse(self, theta, joint_idx):
        angle_deg = np.degrees(theta) % 360  # 规范化角度到0-360

        # 关节1特殊处理
        if joint_idx == 1:
            if angle_deg > 180:
                angle_deg -= 360  # 转换为-180~180范围
            angle_deg = max(angle_deg, 45)  # 强制最小45°

        # 2号关节反向修正
        if joint_idx == 2:
            angle_deg = 180 - angle_deg  # 更合理的反向方式

        pulse = 1500 + (angle_deg - 135) * (2000 / 270)

        # 限制检查
        limits = {1: (600, 2400), 2: (600, 2400), 3: (1100, 2400)}[joint_idx]
        return int(np.clip(pulse, *limits))

    def generate_arm_command(self, target_x_pixel, target_y_pixel):
        x_arm, y_arm = self.image_to_arm_coords(target_x_pixel, target_y_pixel)
        print(f"\n目标坐标: ({target_x_pixel},{target_y_pixel})px → ({x_arm:.1f},{y_arm:.1f})mm")

        solutions = self.inverse_kinematics_all(x_arm, y_arm)
        if not solutions:
            print("逆运动学无有效解")
            return None

        for theta1, theta2, theta3 in solutions:
            pulses = []
            valid = True

            for i, angle in enumerate([theta1, theta2, theta3]):
                pulse = self.angle_to_pulse(angle, i + 1)
                if pulse is None:
                    valid = False
                    break
                pulses.append(pulse)

            if valid:
                return f"#001P{pulses[0]}!#002P{pulses[1]}!#003P{pulses[2]}!"

        print("所有解均不符合关节限制")
        return None
# def compare_parameters(s):
#     for
#         parts = s.split('P')
#         temp_part = parts[0][2:]  # 去掉 "#T"
#         pressure_part, time_part = parts[1].split('T')
#         time_part = time_part.rstrip('！')  # 去掉末尾的 "！"
#         result1 = int(temp_part), int(pressure_part), int(time_part)
#         result2 = '#{}P{}T0000'.format(parts[0][1:], parts[1])
#     return result1,result2


def f(angles):
    controller = ArmController(m, x1, y1)  # m 是主程序中加载的图像

    print(angles)
    for result in angles:
        for key, value in result.items():
            command: str = '#T000P{}T1000!'.format(key)
            print(command)

            command = '#T000P{}T0000!'.format(key)

            print('转动0')
            x, y = value
            comm = controller.generate_arm_command(x, y)
            if comm is None:  # 检查是否返回None
                print(f"无法到达位置 ({x}, {y})，跳过")
                continue

            result = '{' + command + comm + '#005P1500T5000!' + '}'
            print(result)
            print('力臂移动')
            sleep(1)

            c = '{' + command + '#001P1000T1000!#002P1500T1000!#003P1050T1000!#005P1500T0000!' + '}'
            print(c)
            print('抓取物品')
            sleep(1)

            c = '{' + command + '#001P1000T000!#002P1500T0000!#003P1050T0000!#005P1200T1000!' + '}'
            print(c)
            print('放到指定位置')
            sleep(1)

            c = '{' + command + '#001P1400T0000!#002P1700T0000!#003P1400T0000!#005P1500T0000!' + '}'
            print(c)
            print('复位\n')
            sleep(1)
    return None


if __name__ == '__main__':
    # 测试代码
    m = cv2.imread('1.jpg')
    if m is None:
        print("无法加载图像，请检查路径")
        exit(1)
    x1 = 0  # 目标点x坐标
    y1 = 0  # 目标点y坐标
    n = input("请输入要识别的颜色：").strip().lower().split()
    c = Calculate(m, x1, y1)
    angles, mask, display = c.shibie(n)
    controller = ArmController(m, x1, y1)
    f(angles)

    # 显示结果
    cv2.imshow('Detection Result', display)
    cv2.imshow('Color Mask', mask)

    cv2.waitKey(0)
    cv2.destroyAllWindows()
